import gradio as gr
from colorthief import ColorThief
from PIL import Image
import io

# 颜色转换工具
def rgb_to_hex(rgb):
    return "#{:02x}{:02x}{:02x}".format(rgb[0], rgb[1], rgb[2])

# 提取调色板
def extract_palette(image_bytes):
    color_thief = ColorThief(image_bytes)
    palette = color_thief.get_palette(color_count=5)
    hex_palette = [rgb_to_hex(color) for color in palette]
    return palette, hex_palette

# 绘制颜色块
def draw_color_blocks(palette):
    colors = [Image.new("RGB", (100, 100), color=tuple(color)) for color in palette]
    color_blocks = Image.new("RGB", (500, 100))
    for i, block in enumerate(colors):
        color_blocks.paste(block, (i * 100, 0))
    return color_blocks

# 构建颜色详情文本
def build_color_details(palette, hex_palette):
    result = []
    for idx, (color, hex_color) in enumerate(zip(palette, hex_palette)):
        result.append(f"Color {idx + 1}: {color} | {hex_color}")
    return "\n".join(result)

# Gradio 界面构建
def main():
    with gr.Blocks() as app:
        gr.Markdown("## 主色调调色板分析工具")
        gr.Markdown("上传图片，获取其主色调调色板及对应的十六进制颜色代码。")

        with gr.Row():
            image_input = gr.Image(type="pil", label="上传图片")
            output_image = gr.Image(label="主色调调色板")
        
        output_text = gr.Textbox(label="颜色详情")
        submit_button = gr.Button("分析")

        def process(image):
            image_bytes = io.BytesIO()
            image.save(image_bytes, format="PNG")
            image_bytes.seek(0)
            palette, hex_palette = extract_palette(image_bytes)
            color_blocks = draw_color_blocks(palette)
            details = build_color_details(palette, hex_palette)
            return color_blocks, details

        submit_button.click(process, inputs=[image_input], outputs=[output_image, output_text])

    app.launch()

if __name__ == "__main__":
    main()